#! /usr/bin/python3
# -*- coding: UTF-8 -*-
  
'''※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※
File Name: secure.py
Author: GID5564
Description: 非对称加解密
Version: 1.0
Created Time: 11/07/24-18:55:19
※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※'''

#pip install cryptography

from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_OAEP
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.backends import default_backend

import base64
import os

import hashlib
from typing import Union


class Rsa_Pem():
    def __init():
        pass
    def __save_key_to_file(self,key, filename):
        """
        将密钥保存到文件中
        key: RSA密钥对象
        filename: 保存密钥的文件名
        """
        with open(filename, "wb") as f:
            f.write(key.private_bytes(encoding=serialization.Encoding.PEM,format=serialization.PrivateFormat.PKCS8,encryption_algorithm=serialization.NoEncryption()))

    def __load_key_from_file(self,filename):
        """
        从文件中加载密钥
        filename: 包含密钥的文件名
        返回：RSA密钥对象
        """
        with open(filename, "rb") as f:
            return serialization.load_pem_private_key(f.read(),password=None,backend=default_backend())
            
    def save_key(self,pub_key,pri_key):
        self.__save_key_to_file(pub_key,'public_key.pem')
        self.__save_key_to_file(pri_key,'private_key.pem')
        
    def load_key(self,public_pem,private_pem):
        pub_key=self.__load_key_from_file(public_pem)
        pri_key=self.__load_key_from_file(private_pem)
        return pri_key,pub_key


class Rsa_String(Rsa_Pem): #字符串
    def __init__(self,size=2048):
        self.pub_key,self.pri_key=self.__get_key(size)
        
    def __get_key(self,size):
        """
        生成RSA密钥对（公钥和私钥）
        返回：(public_key, private_key)
        """
        key = RSA.generate(size)
        private_key = key.export_key()
        public_key = key.publickey().export_key()
        return public_key, private_key

    def encrypt(self,message,public_key=None):
        """
        使用公钥加密消息
        参数：
            message (str): 要加密的消息
            public_key (str): 用于加密的公钥
        返回：
            str: 加密后的消息（base64编码）
        """
        if not public_key or not message:
            return 0

        rsa_public_key = RSA.import_key(public_key)
        cipher = PKCS1_OAEP.new(rsa_public_key)
        encrypted_message = cipher.encrypt(message.encode())
        return base64.b64encode(encrypted_message).decode()

    def decrypt(self,encrypted_message,private_key=None):
        """
        使用私钥解密消息
        参数：
            encrypted_message (str): 要解密的消息（base64编码）
            pri_key (str): 用于解密的私钥
        返回：
            str: 解密后的消息
        """
        if not private_key or not encrypted_message:
            return 0

        rsa_private_key = RSA.import_key(private_key)
        cipher = PKCS1_OAEP.new(rsa_private_key)
        decoded_message = base64.b64decode(encrypted_message.encode())
        decrypted_message = cipher.decrypt(decoded_message)
        return decrypted_message.decode()


class Rsa_File(Rsa_Pem): #文件
    def __init__(self,size=2048):
        self.pri_key,self.pub_key=self.__get_key(size)
        

    def __get_key(self,size):
        """
        生成RSA密钥对
        返回：(私钥， 公钥)
        """
        private_key = rsa.generate_private_key(public_exponent=65537,key_size=size,backend=default_backend())
        public_key = private_key.public_key()
        return private_key, public_key



    def encrypt_file(self,input_file, output_file, public_key=None):
        """
        使用公钥加密文件
        input_file: 要加密的文件名
        output_file: 加密后的文件名
        public_key: RSA公钥对象
        """
        if not public_key:
            return 0
        try:
            with open(input_file, "rb") as f:
                data = f.read()
            encrypted_data = public_key.encrypt(data,padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()),algorithm=hashes.SHA256(),label=None))
            with open(output_file, "wb") as f:
                f.write(encrypted_data)
        except Exception as e:
            print(e)

    def decrypt_file(self,input_file, output_file, private_key=None):
        """
        使用私钥解密文件
        input_file: 要解密的文件名
        output_file: 解密后的文件名
        private_key: RSA私钥对象
        """
        if not private_key:
            return 0
        try:
            with open(input_file, "rb") as f:
                encrypted_data = f.read()

            decrypted_data = private_key.decrypt(encrypted_data,padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()),algorithm=hashes.SHA256(),label=None))

            with open(output_file, "wb") as f:
                f.write(decrypted_data)
        except Exception as e:
            print(e)

class Hash_Md5():
    def __init__(self):
        pass
    def md5(self,data: Union[str, bytes]) -> str:
        """
        计算给定字符串或字节数据的MD5哈希值。
        参数：
        data (Union[str, bytes]): 要计算哈希值的字符串或字节数据
        返回：
        str: 计算出的MD5哈希值（十六进制表示）。
        """
        # 如果输入是字符串，则将其转换为字节
        if isinstance(data, str):
            data = data.encode('utf-8')
        # 创建一个md5对象
        md5 = hashlib.md5()
        # 更新md5对象的数据
        md5.update(data)
        # 获取十六进制的哈希值
        return md5.hexdigest()
        

    def md5_str(self, text: str):
        md5 = hashlib.md5()
        md5.update(text.encode('utf-8'))
        return md5.hexdigest()

    def md5_file(self, file):
        try:
            # 对文件内容进行MD5哈希
            with open(file, "rb") as f:
                md5_hash = hashlib.md5()
                while chunk := f.read(8192):
                    md5_hash.update(chunk)
                return md5_hash.hexdigest()
        except Exception as e:
            print(e)



if __name__ == "__main__":
    pass







































